Attention(23) CUDA #26466
Merged
Attention(23) CUDA #26466
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This reverts commit 4db63bc.
tianleiwu
reviewed
Jan 9, 2026
tianleiwu
reviewed
Jan 9, 2026
xadupre
reviewed
Jan 13, 2026
tianleiwu
reviewed
Jan 14, 2026
tianleiwu
reviewed
Jan 14, 2026
tianleiwu
previously approved these changes
Jan 14, 2026
tianleiwu
approved these changes
Jan 14, 2026
titaiwangms
added a commit
that referenced
this pull request
Jan 15, 2026
This reverts commit a3e477e.
hariharans29
pushed a commit
that referenced
this pull request
Jan 15, 2026
Reverts #26466 Failing [Attention(23) CUDA (#26466) · a3e477e](https://github.com/microsoft/onnxruntime/actions/runs/21018062539/job/60429722724)
titaiwangms
added a commit
that referenced
this pull request
Jan 15, 2026
This pull request introduces significant improvements and expanded support for multi-head attention kernels in ONNX Runtime, particularly focusing on supporting both 3D (`BSNH`) and 4D (`BNSH`) QKV input formats. The changes enhance flexibility, correctness, and maintainability for attention operations across CPU and CUDA implementations. ### Expanded QKV Input Format Support * Added support for 4D QKV input format (`Q_K_V_BNSH`) in CUDA attention kernels, including proper handling for both cases with and without past/present states, and enforcing that bias is not supported for this format. This includes logic to avoid unnecessary transposes and to write outputs directly when possible. [[1]](diffhunk://#diff-64c7062a412bd7e91378e5c40574de5a1bf63f42ec4cf7d2d23e812fde5bcd11R264-R265) [[2]](diffhunk://#diff-64c7062a412bd7e91378e5c40574de5a1bf63f42ec4cf7d2d23e812fde5bcd11R343-R354) [[3]](diffhunk://#diff-64c7062a412bd7e91378e5c40574de5a1bf63f42ec4cf7d2d23e812fde5bcd11R388-L388) [[4]](diffhunk://#diff-64c7062a412bd7e91378e5c40574de5a1bf63f42ec4cf7d2d23e812fde5bcd11R426-R435) [[5]](diffhunk://#diff-64c7062a412bd7e91378e5c40574de5a1bf63f42ec4cf7d2d23e812fde5bcd11L673-R716) [[6]](diffhunk://#diff-64c7062a412bd7e91378e5c40574de5a1bf63f42ec4cf7d2d23e812fde5bcd11R747-R748) [[7]](diffhunk://#diff-25a30e78aab7a4cdd1d6ba9f3576fc36b79dd3404225d77ea2ee0018490a83eaL775-R791) ### Kernel and Operator Documentation Updates * Updated `OperatorKernels.md` to document the new `Attention` operator inputs and outputs for both 3D and 4D formats, specifying supported tensor types for each input. ### Correctness and Consistency Fixes * Fixed the computation of causal attention indices in CUDA softmax kernels by clarifying and correcting the offset calculation for causal masking. [[1]](diffhunk://#diff-5367f3a93f596de362b09239a92fd1199b3c62fdded9e790810c80526ff9ec9bL168-R168) [[2]](diffhunk://#diff-5367f3a93f596de362b09239a92fd1199b3c62fdded9e790810c80526ff9ec9bL244-R244) [[3]](diffhunk://#diff-5367f3a93f596de362b09239a92fd1199b3c62fdded9e790810c80526ff9ec9bL336-R336) [[4]](diffhunk://#diff-5367f3a93f596de362b09239a92fd1199b3c62fdded9e790810c80526ff9ec9bL442-R442) * Updated workspace allocation logic for QKV preparation to ensure correct workspace usage for new formats. ### Attention Parameter and Helper Refactoring * Added `is_output_bnsh` field to `AttentionParameters` to indicate output format and updated logic to use this for output placement and transposition decisions. [[1]](diffhunk://#diff-e742290164e1e1fa0152840db2a1b83354e153153df19a2762b58655e49b7f9bR37) [[2]](diffhunk://#diff-25a30e78aab7a4cdd1d6ba9f3576fc36b79dd3404225d77ea2ee0018490a83eaL775-R791) * Refactored CPU attention implementation to use the new `attention_helper` namespace for output mode enums and output shape computation, improving code clarity and maintainability. [[1]](diffhunk://#diff-e692b5c865c4874e51982867901cd514e68cf38dd435c00fe505f34f93956fe7R5) [[2]](diffhunk://#diff-e692b5c865c4874e51982867901cd514e68cf38dd435c00fe505f34f93956fe7L118-R125) [[3]](diffhunk://#diff-e692b5c865c4874e51982867901cd514e68cf38dd435c00fe505f34f93956fe7L143-R149) ### Minor Cleanups * Removed outdated asserts and improved debug output strings for QKV preparation functions to clarify format and state handling. [[1]](diffhunk://#diff-64c7062a412bd7e91378e5c40574de5a1bf63f42ec4cf7d2d23e812fde5bcd11L254) [[2]](diffhunk://#diff-64c7062a412bd7e91378e5c40574de5a1bf63f42ec4cf7d2d23e812fde5bcd11L363) [[3]](diffhunk://#diff-64c7062a412bd7e91378e5c40574de5a1bf63f42ec4cf7d2d23e812fde5bcd11L673-R716) These changes collectively improve the flexibility, correctness, and maintainability of attention kernel implementations in ONNX Runtime, especially for advanced transformer models and large language model workloads. **NOT supported in this PR** - Boolean mask - GQA - Softcap - Softmax precision - qk_output_mode other than -1 and 0
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This pull request introduces significant improvements and expanded support for multi-head attention kernels in ONNX Runtime, particularly focusing on supporting both 3D (
BSNH) and 4D (BNSH) QKV input formats. The changes enhance flexibility, correctness, and maintainability for attention operations across CPU and CUDA implementations.Expanded QKV Input Format Support
Q_K_V_BNSH) in CUDA attention kernels, including proper handling for both cases with and without past/present states, and enforcing that bias is not supported for this format. This includes logic to avoid unnecessary transposes and to write outputs directly when possible. [1] [2] [3] [4] [5] [6] [7]Kernel and Operator Documentation Updates
OperatorKernels.mdto document the newAttentionoperator inputs and outputs for both 3D and 4D formats, specifying supported tensor types for each input.Correctness and Consistency Fixes
Attention Parameter and Helper Refactoring
is_output_bnshfield toAttentionParametersto indicate output format and updated logic to use this for output placement and transposition decisions. [1] [2]attention_helpernamespace for output mode enums and output shape computation, improving code clarity and maintainability. [1] [2] [3]Minor Cleanups
These changes collectively improve the flexibility, correctness, and maintainability of attention kernel implementations in ONNX Runtime, especially for advanced transformer models and large language model workloads.
NOT supported in this PR